# _*_coding:utf-8_*_
# author:buka
import glob
import itertools

import cv2
import numpy as np
import tensorflow as tf
from model.crnn.Model import get_Model
from model.crnn.parameter import letters

graph = tf.get_default_graph()
with graph.as_default():
    crnn = get_Model(False)
    crnn.load_weights('./model/crnn/LSTM+BN5--383--0.000.hdf5')

def show(pan, img):
    cv2.destroyWindow(pan)
    cv2.imshow(str(pan), cv2.resize(img, (0, 0), fx = 0.5, fy = 0.5))
    cv2.waitKey()

def locate(img):
    grad = cv2.morphologyEx(img, cv2.MORPH_GRADIENT, np.ones((7, 7)))
    dltd = cv2.dilate(grad, np.ones((17, 17), np.uint8))
    show('dltd', dltd)

    cnts = cv2.findContours(dltd, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
    cnts = sorted(cnts, key = lambda cnt: cv2.contourArea(cnt), reverse = True)[:5]
    cntd = cv2.drawContours(cv2.cvtColor(img, cv2.COLOR_GRAY2BGR), cnts, -1, (255, 0, 0), 2)
    show('cntd', cntd)

    cntrs = []
    for cnt in cnts:
        cntr = cv2.minAreaRect(cnt)
        if cntr[1][0] / cntr[1][1] > 1:
            cntrs.append(cntr)
        else:
            cntr = (cntr[0], (cntr[1][1], cntr[1][0]), cntr[2] + 90)
            cntrs.append(cntr)

    cntrws = list(map(lambda cntr: cntr[0][1] * cntr[1][0] / cntr[1][1], cntrs))
    cnti = int(np.argmax(cntrws))
    cntd = cv2.drawContours(cv2.cvtColor(img, cv2.COLOR_GRAY2BGR), [cnts[cnti]], -1, (255, 0, 0), 2)
    show('cntd', cntd)

    box = cv2.boxPoints(cntrs[cnti])
    cntbd = cv2.drawContours(cv2.cvtColor(img, cv2.COLOR_GRAY2BGR), [np.int0(box)], -1, (0, 0, 255), 2)
    show('cntd', cntbd)

    img_w, img_h = 256, 32
    boxd = np.float32([[0, img_h], [0, 0], [img_w, 0], [img_w, img_h]])
    M = cv2.getPerspectivM = cv2.getPerspectiveTransform(box, boxd)
    boxd = cv2.warpPerspective(img, M, (img_w, img_h))
    show('boxd', boxd)
    return boxd

def reco(img_pred):
    img_pred = (img_pred / 255.0) * 2.0 - 1.0
    img_pred = img_pred.T
    img_pred = np.expand_dims(img_pred, axis = -1)
    img_pred = np.expand_dims(img_pred, axis = 0)

    with graph.as_default():
        net_out_value = crnn.predict(img_pred)
    pred_best = list(np.argmax(net_out_value[0, 2:], axis = 1))  # get max index -> len = 32
    pred_best = [k for k, g in itertools.groupby(pred_best)]  # remove overlap value
    pred_text = ''
    for i in pred_best:
        if i < len(letters):
            pred_text += letters[i]
    return pred_text

import time
for file in glob.glob('./img/*.bmp'):
    img = cv2.imread(file, 0)
    t1 = time.time()
    img_pred = locate(img)
    t2 = time.time()
    pred_text = reco(img_pred)
    print(file, pred_text)
    print(t2-t1,time.time()-t2)
